# Adopted from https://github.com/lm-sys/FastChat. Below is the original copyright:
# Adopted from tatsu-lab@stanford_alpaca. Below is the original copyright:
#    Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.

import os
import copy
from dataclasses import dataclass, field
import random
import json
import logging
import pathlib
from typing import Dict, Optional, Sequence

import torch

import transformers
from transformers import GPTNeoXTokenizerFast
import tokenizers

from llava.constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
from torch.utils.data import Dataset
from llava.train.llava_trainer import LLaVATrainer

from llava import conversation as conversation_lib
from llava.model import *
from llava.mm_utils import tokenizer_image_token

from PIL import Image


local_rank = None


def rank0_print(*args):
    if local_rank == 0:
        print(*args)


from packaging import version
IS_TOKENIZER_GREATER_THAN_0_14 = version.parse(tokenizers.__version__) >= version.parse('0.14')


@dataclass
class ModelArguments:
    model_name_or_path: Optional[str] = field(default="facebook/opt-125m")
    version: Optional[str] = field(default="v0")
    freeze_backbone: bool = field(default=False)
    tune_mm_mlp_adapter: bool = field(default=False)
    tune_vision_tower: bool = field(default=False)
    vision_tower: Optional[str] = field(default=None)
    mm_vision_select_layer: Optional[int] = field(default=-1)   # default to the last layer
    pretrain_mm_mlp_adapter: Optional[str] = field(default=None)
    mm_projector_type: Optional[str] = field(default='linear')
    mm_use_im_start_end: bool = field(default=False)
    mm_use_im_patch_token: bool = field(default=True)
    mm_patch_merge_type: Optional[str] = field(default='flat')
    mm_vision_select_feature: Optional[str] = field(default="patch")


@dataclass
class DataArguments:
    data_path: str = field(default=None,
                           metadata={"help": "Path to the training data."})
    coco_caption_prompt_file: str = field(default=None)
    lazy_preprocess: bool = False
    is_multimodal: bool = False
    image_folder: Optional[str] = field(default=None)
    image_aspect_ratio: str = 'square'


@dataclass
class TrainingArguments(transformers.TrainingArguments):
    cache_dir: Optional[str] = field(default=None)
    optim: str = field(default="adamw_torch")
    remove_unused_columns: bool = field(default=False)
    freeze_mm_mlp_adapter: bool = field(default=False)
    mpt_attn_impl: Optional[str] = field(default="triton")
    model_max_length: int = field(
        default=512,
        metadata={
            "help":
            "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
        },
    )
    double_quant: bool = field(
        default=True,
        metadata={"help": "Compress the quantization statistics through double quantization."}
    )
    quant_type: str = field(
        default="nf4",
        metadata={"help": "Quantization data type to use. Should be one of `fp4` or `nf4`."}
    )
    bits: int = field(
        default=16,
        metadata={"help": "How many bits to use."}
    )
    lora_enable: bool = False
    lora_r: int = 64
    lora_alpha: int = 16
    lora_dropout: float = 0.05
    lora_weight_path: str = ""
    lora_bias: str = "none"
    mm_projector_lr: Optional[float] = None
    mm_projector_wd: Optional[float] = None
    vision_tower_lr: Optional[float] = None
    vision_tower_lldr: Optional[float] = None
    group_by_modality_length: bool = field(default=False)
    coco_caption_sft: bool = field(default=False)
    # My additions: Noised embeddings finetuning (NEFTune) for MLLMs
    neftune_alpha: Optional[float] = None


def maybe_zero_3(param, ignore_status=False, name=None):
    from deepspeed import zero
    from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
    if hasattr(param, "ds_id"):
        if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
            if not ignore_status:
                logging.warning(f"{name}: param.ds_status != ZeroParamStatus.NOT_AVAILABLE: {param.ds_status}")
        with zero.GatheredParameters([param]):
            param = param.data.detach().cpu().clone()
    else:
        param = param.detach().cpu().clone()
    return param


# Borrowed from peft.utils.get_peft_model_state_dict
def get_peft_state_maybe_zero_3(named_params, bias):
    if bias == "none":
        to_return = {k: t for k, t in named_params if "lora_" in k}
    elif bias == "all":
        to_return = {k: t for k, t in named_params if "lora_" in k or "bias" in k}
    elif bias == "lora_only":
        to_return = {}
        maybe_lora_bias = {}
        lora_bias_names = set()
        for k, t in named_params:
            if "lora_" in k:
                to_return[k] = t
                bias_name = k.split("lora_")[0] + "bias"
                lora_bias_names.add(bias_name)
            elif "bias" in k:
                maybe_lora_bias[k] = t
        for k, t in maybe_lora_bias:
            if bias_name in lora_bias_names:
                to_return[bias_name] = t
    else:
        raise NotImplementedError
    to_return = {k: maybe_zero_3(v, ignore_status=True) for k, v in to_return.items()}
    return to_return


def get_peft_state_non_lora_maybe_zero_3(named_params, require_grad_only=True):
    to_return = {k: t for k, t in named_params if "lora_" not in k}
    if require_grad_only:
        to_return = {k: t for k, t in to_return.items() if t.requires_grad}
    to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()}
    return to_return


def get_mm_adapter_state_maybe_zero_3(named_params, keys_to_match):
    to_return = {k: t for k, t in named_params if any(key_match in k for key_match in keys_to_match)}
    to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()}
    return to_return


# My additions, for CLIP-ViT saving when sft vision tower
def get_vision_tower_state_maybe_zero_3(named_params, keys_to_match=["vision_model"]):
    to_return = {k: t for k, t in named_params if any(key_match in k for key_match in keys_to_match)}
    to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()}
    return to_return


def find_all_linear_names(model):
    cls = torch.nn.Linear
    lora_module_names = set()
    multimodal_keywords = ['mm_projector', 'vision_tower', 'vision_resampler']
    for name, module in model.named_modules():
        if any(mm_keyword in name for mm_keyword in multimodal_keywords):
            continue
        if isinstance(module, cls):
            names = name.split('.')
            lora_module_names.add(names[0] if len(names) == 1 else names[-1])

    if 'lm_head' in lora_module_names: # needed for 16-bit
        lora_module_names.remove('lm_head')
    return list(lora_module_names)


# My additions, for CLIP ViT saving when it is finetuned
def safe_save_vision_tower_for_hf_trainer(vision_tower_folder: str, clip_vit_model: transformers.CLIPVisionModel, 
                                          vision_tower_weights_to_save:Dict, clip_image_processor: transformers.CLIPImageProcessor):
    if vision_tower_folder is not None:  # true means saving CLIP ViT model
        if (clip_vit_model is None) or (vision_tower_weights_to_save is None) or (clip_image_processor is None):
            assert 0
        clip_image_processor.save_pretrained(vision_tower_folder)  # saving hf CLIPImageProcessor instance
        clip_vit_model.config.save_pretrained(vision_tower_folder)  # saving hf CLIPVisionModel instance
        torch.save(vision_tower_weights_to_save, os.path.join(vision_tower_folder, "pytorch_model.bin"))  # saving CLIP-ViT weights
        del vision_tower_weights_to_save


def safe_save_model_for_hf_trainer(trainer: transformers.Trainer,
                                   output_dir: str):
    """Collects the state dict and dump to disk."""

    if getattr(trainer.args, "tune_vision_tower", False):
        vision_tower_name = getattr(trainer.args, "vision_tower_name_or_path", "openai/clip-vit-large-patch14-336").split("/")[-1]
        vision_tower_folder = os.path.join(output_dir, vision_tower_name)

        clip_vit_model = getattr(trainer.model.get_vision_tower(), "vision_tower", None)
        vision_tower_weights_to_save = get_vision_tower_state_maybe_zero_3(clip_vit_model.named_parameters())
        clip_image_processor = getattr(trainer.model.get_vision_tower(), "image_processor")
    else:
        vision_tower_folder = clip_vit_model = vision_tower_weights_to_save = clip_image_processor = None

    if getattr(trainer.args, "tune_mm_mlp_adapter", False):
        # Only save Adapter
        mm_projector_keys_to_match = ['mm_projector']
        if getattr(trainer.args, "use_im_start_end", False):
            mm_projector_keys_to_match.extend(['embed_tokens', 'embed_in', 'embeddings'])

        mm_projector_weights_to_save = get_mm_adapter_state_maybe_zero_3(trainer.model.named_parameters(), mm_projector_keys_to_match)
        trainer.model.config.save_pretrained(output_dir)

        current_folder = output_dir.split('/')[-1]
        parent_folder = os.path.dirname(output_dir)
        if trainer.args.local_rank == 0 or trainer.args.local_rank == -1:
            if current_folder.startswith('checkpoint-'):
                mm_projector_folder = os.path.join(parent_folder, "mm_projector")
                os.makedirs(mm_projector_folder, exist_ok=True)
                torch.save(mm_projector_weights_to_save, os.path.join(mm_projector_folder, f'{current_folder}.bin'))
            else:
                torch.save(mm_projector_weights_to_save, os.path.join(output_dir, f'mm_projector.bin'))
            # My additions, for CLIP ViT saving when it is finetuned
            safe_save_vision_tower_for_hf_trainer(vision_tower_folder, clip_vit_model, vision_tower_weights_to_save, clip_image_processor)
        return

    if trainer.deepspeed:
        torch.cuda.synchronize()
        trainer.save_model(output_dir)
        # My additions, for CLIP ViT saving 
        if trainer.args.should_save:
            safe_save_vision_tower_for_hf_trainer(vision_tower_folder, clip_vit_model, vision_tower_weights_to_save, clip_image_processor)
        return

    state_dict = trainer.model.state_dict()
    if trainer.args.should_save:
        cpu_state_dict = {
            key: value.cpu()
            for key, value in state_dict.items()
        }
        del state_dict
        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa
        # My additions, for CLIP ViT saving
        safe_save_vision_tower_for_hf_trainer(vision_tower_folder, clip_vit_model, vision_tower_weights_to_save, clip_image_processor)


def smart_tokenizer_and_embedding_resize(
    special_tokens_dict: Dict,
    tokenizer: transformers.PreTrainedTokenizer,
    model: transformers.PreTrainedModel,
):
    """Resize tokenizer and embedding.

    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
    """
    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
    model.resize_token_embeddings(len(tokenizer))

    if num_new_tokens > 0:
        input_embeddings = model.get_input_embeddings().weight.data
        output_embeddings = model.get_output_embeddings().weight.data

        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
            dim=0, keepdim=True)
        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
            dim=0, keepdim=True)

        input_embeddings[-num_new_tokens:] = input_embeddings_avg
        output_embeddings[-num_new_tokens:] = output_embeddings_avg


def _tokenize_fn(strings: Sequence[str],
                 tokenizer: transformers.PreTrainedTokenizer) -> Dict:
    """Tokenize a list of strings."""
    tokenized_list = [
        tokenizer(
            text,
            return_tensors="pt",
            padding="longest",
            max_length=tokenizer.model_max_length,
            truncation=True,
        ) for text in strings
    ]
    input_ids = labels = [
        tokenized.input_ids[0] for tokenized in tokenized_list
    ]
    input_ids_lens = labels_lens = [
        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item()
        for tokenized in tokenized_list
    ]
    return dict(
        input_ids=input_ids,
        labels=labels,
        input_ids_lens=input_ids_lens,
        labels_lens=labels_lens,
    )


def _mask_targets(target, tokenized_lens, speakers):
    # cur_idx = 0
    cur_idx = tokenized_lens[0]
    tokenized_lens = tokenized_lens[1:]
    target[:cur_idx] = IGNORE_INDEX
    for tokenized_len, speaker in zip(tokenized_lens, speakers):
        if speaker == "human":
            target[cur_idx+2:cur_idx + tokenized_len] = IGNORE_INDEX
        cur_idx += tokenized_len


def _add_speaker_and_signal(header, source, get_conversation=True):
    """Add speaker and start/end signal on each round."""
    BEGIN_SIGNAL = "### "
    END_SIGNAL = "\n"
    conversation = header
    for sentence in source:
        from_str = sentence["from"]
        if from_str.lower() == "human":
            from_str = conversation_lib.default_conversation.roles[0]
        elif from_str.lower() == "gpt":
            from_str = conversation_lib.default_conversation.roles[1]
        else:
            from_str = 'unknown'
        sentence["value"] = (BEGIN_SIGNAL + from_str + ": " +
                             sentence["value"] + END_SIGNAL)
        if get_conversation:
            conversation += sentence["value"]
    conversation += BEGIN_SIGNAL
    return conversation


def preprocess_multimodal(
    # My additions
    # preprocess_multimodal这个函数也只是处理文本而已,只不过这是对应多模态对话的文本
    # sources中并没有图像的Tensor,全是字符串
    # 如果是multimodal的输入,那sources不应该是Sequence[str],其实际的格式应该如下:Sequence[List[Dict]]
    # Sequence[List[Dict]]中的每一个List元素对应着一次多模态对话,而List中的每一个Dict代表着本次多模态对话
    # 中的一次提问或者回答,Dict中的键一般情况为'from'和'value'
    sources: Sequence[str],
    data_args: DataArguments
) -> Dict:
    is_multimodal = data_args.is_multimodal
    if not is_multimodal:
        return sources

    for source in sources:
        for sentence in source:
            if DEFAULT_IMAGE_TOKEN in sentence['value']:
                sentence['value'] = sentence['value'].replace(DEFAULT_IMAGE_TOKEN, '').strip()
                sentence['value'] = DEFAULT_IMAGE_TOKEN + '\n' + sentence['value']
                sentence['value'] = sentence['value'].strip()
                if "mmtag" in conversation_lib.default_conversation.version:
                    sentence['value'] = sentence['value'].replace(DEFAULT_IMAGE_TOKEN, '<Image>' + DEFAULT_IMAGE_TOKEN + '</Image>')
            replace_token = DEFAULT_IMAGE_TOKEN
            if data_args.mm_use_im_start_end:
                replace_token = DEFAULT_IM_START_TOKEN + replace_token + DEFAULT_IM_END_TOKEN
            sentence["value"] = sentence["value"].replace(DEFAULT_IMAGE_TOKEN, replace_token)

    return sources


def preprocess_llama_2(
    sources,
    tokenizer: transformers.PreTrainedTokenizer,
    has_image: bool = False
) -> Dict:
    conv = conversation_lib.default_conversation.copy()
    roles = {"human": conv.roles[0], "gpt": conv.roles[1]}

    # Apply prompt templates
    conversations = []
    for i, source in enumerate(sources):
        if roles[source[0]["from"]] != conv.roles[0]:
            # Skip the first one if it is not from human
            source = source[1:]

        conv.messages = []
        for j, sentence in enumerate(source):
            role = roles[sentence["from"]]
            assert role == conv.roles[j % 2], f"{i}"
            conv.append_message(role, sentence["value"])
        conversations.append(conv.get_prompt())

    # Tokenize conversations

    if has_image:
        input_ids = torch.stack([tokenizer_image_token(prompt, tokenizer, return_tensors='pt') for prompt in conversations], dim=0)
    else:
        input_ids = tokenizer(
            conversations,
            return_tensors="pt",
            padding="longest",
            max_length=tokenizer.model_max_length,
            truncation=True,
        ).input_ids

    targets = input_ids.clone()

    assert conv.sep_style == conversation_lib.SeparatorStyle.LLAMA_2

    # Mask targets
    sep = "[/INST] "
    for conversation, target in zip(conversations, targets):
        total_len = int(target.ne(tokenizer.pad_token_id).sum())

        rounds = conversation.split(conv.sep2)
        cur_len = 1
        target[:cur_len] = IGNORE_INDEX
        for i, rou in enumerate(rounds):
            if rou == "":
                break

            parts = rou.split(sep)
            if len(parts) != 2:
                break
            parts[0] += sep

            if has_image:
                round_len = len(tokenizer_image_token(rou, tokenizer))
                instruction_len = len(tokenizer_image_token(parts[0], tokenizer)) - 2
            else:
                round_len = len(tokenizer(rou).input_ids)
                instruction_len = len(tokenizer(parts[0]).input_ids) - 2

            target[cur_len : cur_len + instruction_len] = IGNORE_INDEX

            cur_len += round_len
        target[cur_len:] = IGNORE_INDEX

        if cur_len < tokenizer.model_max_length:
            if cur_len != total_len:
                target[:] = IGNORE_INDEX
                print(
                    f"WARNING: tokenization mismatch: {cur_len} vs. {total_len}."
                    f" (ignored)"
                )

    return dict(
        input_ids=input_ids,
        labels=targets,
    )


def preprocess_v1(
    sources,
    tokenizer: transformers.PreTrainedTokenizer,
    has_image: bool = False
) -> Dict:
    conv = conversation_lib.default_conversation.copy()
    roles = {"human": conv.roles[0], "gpt": conv.roles[1]}

    # Apply prompt templates
    conversations = []
    for i, source in enumerate(sources):
        if roles[source[0]["from"]] != conv.roles[0]:
            # Skip the first one if it is not from human
            source = source[1:]

        conv.messages = []
        for j, sentence in enumerate(source):
            role = roles[sentence["from"]]
            assert role == conv.roles[j % 2], f"{i}"
            conv.append_message(role, sentence["value"])
        conversations.append(conv.get_prompt())

    # Tokenize conversations

    if has_image:
        input_ids = torch.stack([tokenizer_image_token(prompt, tokenizer, return_tensors='pt') for prompt in conversations], dim=0)
    else:
        input_ids = tokenizer(
            conversations,
            return_tensors="pt",
            padding="longest",
            max_length=tokenizer.model_max_length,
            truncation=True,
        ).input_ids

    targets = input_ids.clone()

    assert conv.sep_style == conversation_lib.SeparatorStyle.TWO

    # Mask targets
    sep = conv.sep + conv.roles[1] + ": "
    for conversation, target in zip(conversations, targets):
        total_len = int(target.ne(tokenizer.pad_token_id).sum())

        rounds = conversation.split(conv.sep2)
        cur_len = 1
        target[:cur_len] = IGNORE_INDEX
        for i, rou in enumerate(rounds):
            if rou == "":
                break

            parts = rou.split(sep)
            if len(parts) != 2:
                break
            parts[0] += sep

            if has_image:
                round_len = len(tokenizer_image_token(rou, tokenizer))
                instruction_len = len(tokenizer_image_token(parts[0], tokenizer)) - 2
            else:
                round_len = len(tokenizer(rou).input_ids)
                instruction_len = len(tokenizer(parts[0]).input_ids) - 2

            if i != 0 and not tokenizer.legacy and IS_TOKENIZER_GREATER_THAN_0_14:
                round_len -= 1
                instruction_len -= 1

            target[cur_len : cur_len + instruction_len] = IGNORE_INDEX

            cur_len += round_len
        target[cur_len:] = IGNORE_INDEX

        if cur_len < tokenizer.model_max_length:
            if cur_len != total_len:
                target[:] = IGNORE_INDEX
                print(
                    f"WARNING: tokenization mismatch: {cur_len} vs. {total_len}."
                    f" (ignored)"
                )

    return dict(
        input_ids=input_ids,
        labels=targets,
    )


# My additions
def preprocess_mpt(
    sources,
    tokenizer: transformers.PreTrainedTokenizer,
    has_image: bool = False,
    gptneox_tokenizer: bool = False,
    coco_caption_sft: bool = False,
    coco_caption_prompts: Sequence = []
) -> Dict:
    conv = conversation_lib.default_conversation.copy()
    roles = {"human": conv.roles[0], "gpt": conv.roles[1]}

    # Apply prompt templates
    conversations = []
    for i, source in enumerate(sources):
        if roles[source[0]["from"]] != conv.roles[0]:
            # Skip the first one if it is not from human
            source = source[1:]

        conv.messages = []
        for j, sentence in enumerate(source):
            role = roles[sentence["from"]]
            assert role == conv.roles[j % 2], f"{i}"
            if coco_caption_sft and sentence["value"] in (DEFAULT_IMAGE_TOKEN, DEFAULT_IMAGE_TOKEN + "\n") and sentence["from"] == "human":
                text_prompt = random.choice(coco_caption_prompts)
                conv.append_message(role, "".join([DEFAULT_IMAGE_TOKEN, "\n", text_prompt]))
            else:
                conv.append_message(role, sentence["value"])
        # My additions
        # conv.get_prompt()返回的是一个str
        conversations.append(conv.get_prompt())

    # Tokenize conversations

    if has_image:
        # My additions
        # input_ids是一个torch.Tensor, input_ids.shape: (1, token_length)
        # input_ids的第一个维度为1,因为目前只支持处理单个的多轮对话的样本
        input_ids = torch.stack([tokenizer_image_token(prompt, tokenizer, return_tensors='pt') for prompt in conversations], dim=0)
    else:
        input_ids = tokenizer(
            conversations,
            return_tensors="pt",
            padding="longest",
            max_length=tokenizer.model_max_length,
            truncation=True,
        ).input_ids

    targets = input_ids.clone()
    assert conv.sep_style == conversation_lib.SeparatorStyle.MPT

    # Mask targets
    sep = conv.sep + conv.roles[1]
    for conversation, target in zip(conversations, targets):
        total_len = int(target.ne(tokenizer.pad_token_id).sum())

        # My additions
        # 在MPT/Mamba的情况下,它们的tokenizer都是GPTNeoXTokenizer,conv.sep是'<|im_end|>'
        # 当将conversation字符串按照conv.sep分开时,由于batch内的对话字符串中,
        # AI的回答一定是以'<|im_end|>'结尾,那么rounds的最后一个元素,即rounds[-1]一定是空字符串''
        # 所以当conv_idx循环到了rounds的最后一个元素的索引,即空字符串''的索引时,
        # conv.sep.join(rounds[conv_idx: conv_idx + 2])一定也是空字符串,
        # 即re_rounds的最后一个元素也一定是空字符串''
        # 所以下面的for loop中,有`if rou == ""`就break的条件判断
        rounds = conversation.split(conv.sep)
        re_rounds = [conv.sep.join(rounds[:3])] # system + user + gpt
        for conv_idx in range(3, len(rounds), 2):
            re_rounds.append(conv.sep.join(rounds[conv_idx:conv_idx+2]))    # user + gpt
        cur_len = 0
        target[:cur_len] = IGNORE_INDEX
        # My additions
        # rou: "<|im_start|>user\n`user_input_str`<|im_end|><|im_start|>assistant\n`assistant_output_str`"
        for i, rou in enumerate(re_rounds):
            if rou == "":
                break

            parts = rou.split(sep)
            if len(parts) != 2:
                break
            # My additions
            # For MPT LLM, sep is `<|im_end|><|im_start|>assistant\n`
            parts[0] += sep

            if has_image:
                round_len = len(tokenizer_image_token(rou, tokenizer))
                # My additions
                # `-1` means LLaVA-MPT-Chat wants the `\n` in `<|im_end|><|im_start|>assistant\n` to be predicted
                # `-1` or `-2` in other conversation templates also have the same meaning
                instruction_len = len(tokenizer_image_token(parts[0], tokenizer)) - 1
            else:
                round_len = len(tokenizer(rou).input_ids)
                instruction_len = len(tokenizer(parts[0]).input_ids) - 1

            if i != 0 and getattr(tokenizer, 'legacy', False) and IS_TOKENIZER_GREATER_THAN_0_14:
                round_len += 1
                instruction_len += 1

            # My additions
            # for GPTNeoXTokenizerFast, it uses eos_token, i.e. adding token `<|im_end|>` at the end of generation,
            # but the str `rou` in List[str] `re_rounds` doesn't have the `<|im_end|>` token in the end,
            # so the variable `round_len` must add one for the lost `<|im_end|>` token
            if gptneox_tokenizer:
                round_len += 1

            target[cur_len : cur_len + instruction_len] = IGNORE_INDEX

            cur_len += round_len
        target[cur_len:] = IGNORE_INDEX

        if cur_len < tokenizer.model_max_length:
            if cur_len != total_len:
                target[:] = IGNORE_INDEX
                print(
                    f"WARNING: tokenization mismatch: {cur_len} vs. {total_len}."
                    f" (ignored)"
                )

    return dict(
        input_ids=input_ids,
        labels=targets,
    )


def preprocess_plain(
    sources: Sequence[str],
    tokenizer: transformers.PreTrainedTokenizer,
    gptneox_tokenizer: bool = False
) -> Dict:
    # add end signal and concatenate together
    conversations = []
    for source in sources:
        assert len(source) == 2
        assert DEFAULT_IMAGE_TOKEN in source[0]['value']
        # The following comments and `DEFAULT_IMAGE_TOKEN + '\nA picture of '` is used for:
        # using plain mode conversation to train coco karpathy split caption
        # alternative prompt for coco caption CIDEr benchmark
        # source[0]['value'] = DEFAULT_IMAGE_TOKEN + '\nA picture of '
        source[0]['value'] = DEFAULT_IMAGE_TOKEN
        conversation = source[0]['value'] + source[1]['value'] + conversation_lib.default_conversation.sep
        conversations.append(conversation)
    # tokenize conversations
    input_ids = [tokenizer_image_token(prompt, tokenizer, return_tensors='pt') for prompt in conversations]
    targets = copy.deepcopy(input_ids)
    for target, source in zip(targets, sources):
        # My additions for debugging
        if not gptneox_tokenizer:
            tokenized_len = len(tokenizer_image_token(source[0]['value'], tokenizer))
        else:
            # The following comments and `source[0]['value'][: -1]` is used for:
            # using plain mode conversation to train coco karpathy split caption
            # These additions are used for Mamba-VL training during coco karpathy split training
            # When training on coco karpathy split training, the prompt text, 
            # i.e. variable `source[0]['value']` is `<image>\nA picture of `, and the tokenizer of Mamba
            # i.e. an instance of GPTNeoXTokenizerFast, and the tokenizer tokenizes `\nA picture of ` into
            # `['Ċ', 'A', 'Ġpicture', 'Ġof', 'Ġ']`, in which 'Ċ' means '\n' and 'Ġ' means ' ',
            # in fact we want the model to predict the first token of image caption based on 'Ġof', but not 'Ġ',
            # so using `source[0]['value'][: -1]` instead. For example, if the first word in the image caption
            # is 'A', we want the model to predict 'ĠA' based on 'Ġof'
            # What's more, the prompt text for COCO karpathy split caption should be changed to 
            # '<image>\nA picture of', but not '<image>\nA picture of '
            # tokenized_len = len(tokenizer_image_token(source[0]['value'][: -1], tokenizer))
            tokenized_len = len(tokenizer_image_token(source[0]["value"], tokenizer))
        target[:tokenized_len] = IGNORE_INDEX

    return dict(input_ids=input_ids, labels=targets)


# My additions
def preprocess(
    sources: Sequence[str],
    tokenizer: transformers.PreTrainedTokenizer,
    has_image: bool = False,
    gptneox_tokenizer: bool = False,
    coco_caption_sft: bool = False,
    coco_caption_prompts: Sequence = []
) -> Dict:
    """
    Given a list of sources, each is a conversation list. This transform:
    1. Add signal '### ' at the beginning each sentence, with end signal '\n';
    2. Concatenate conversations together;
    3. Tokenize the concatenated conversation;
    4. Make a deepcopy as the target. Mask human words with IGNORE_INDEX.
    gptneox_tokenizer(bool): default False, whether the tokenizer is an instance of GPTNeoXTokenizerFast,
    GPTNeoXTokenizer is used for MPT LLM and Mamba S6 series models.
    """
    if conversation_lib.default_conversation.sep_style == conversation_lib.SeparatorStyle.PLAIN:
        return preprocess_plain(sources, tokenizer, gptneox_tokenizer=gptneox_tokenizer)
    if conversation_lib.default_conversation.sep_style == conversation_lib.SeparatorStyle.LLAMA_2:
        return preprocess_llama_2(sources, tokenizer, has_image=has_image)
    if conversation_lib.default_conversation.version.startswith("v1"):
        return preprocess_v1(sources, tokenizer, has_image=has_image)
    if conversation_lib.default_conversation.version == "mpt":
        return preprocess_mpt(sources, tokenizer, has_image=has_image, gptneox_tokenizer=gptneox_tokenizer, 
                              coco_caption_sft=coco_caption_sft, coco_caption_prompts=coco_caption_prompts)
    # add end signal and concatenate together
    conversations = []
    for source in sources:
        header = f"{conversation_lib.default_conversation.system}\n\n"
        conversation = _add_speaker_and_signal(header, source)
        conversations.append(conversation)
    # tokenize conversations
    def get_tokenize_len(prompts):
        return [len(tokenizer_image_token(prompt, tokenizer)) for prompt in prompts]

    if has_image:
        input_ids = [tokenizer_image_token(prompt, tokenizer, return_tensors='pt') for prompt in conversations]
    else:
        conversations_tokenized = _tokenize_fn(conversations, tokenizer)
        input_ids = conversations_tokenized["input_ids"]

    targets = copy.deepcopy(input_ids)
    for target, source in zip(targets, sources):
        if has_image:
            tokenized_lens = get_tokenize_len([header] + [s["value"] for s in source])
        else:
            tokenized_lens = _tokenize_fn([header] + [s["value"] for s in source], tokenizer)["input_ids_lens"]
        speakers = [sentence["from"] for sentence in source]
        _mask_targets(target, tokenized_lens, speakers)

    return dict(input_ids=input_ids, labels=targets)


class LazySupervisedDataset(Dataset):
    """Dataset for supervised fine-tuning."""

    def __init__(self, data_path: str,
                 tokenizer: transformers.PreTrainedTokenizer,
                 data_args: DataArguments, coco_caption_sft: bool = False):
        super(LazySupervisedDataset, self).__init__()
        list_data_dict = json.load(open(data_path, "r"))
        if coco_caption_sft:
            assert pathlib.Path(data_args.coco_caption_prompt_file).exists()
            coco_caption_prompts = open(data_args.coco_caption_prompt_file, "r").readlines()
            self.coco_caption_prompts = [p.strip() for p in coco_caption_prompts]
        self.coco_caption_sft = coco_caption_sft

        rank0_print("Formatting inputs...Skip in lazy mode")
        self.tokenizer = tokenizer
        # My additions
        if isinstance(tokenizer, GPTNeoXTokenizerFast):
            self.gptneox_tokenizer = True
        else:
            self.gptneox_tokenizer = False
        self.list_data_dict = list_data_dict
        self.data_args = data_args

    def __len__(self):
        return len(self.list_data_dict)

    @property
    def lengths(self):
        length_list = []
        for sample in self.list_data_dict:
            img_tokens = 128 if 'image' in sample else 0
            length_list.append(sum(len(conv['value'].split()) for conv in sample['conversations']) + img_tokens)
        return length_list

    @property
    def modality_lengths(self):
        length_list = []
        for sample in self.list_data_dict:
            cur_len = sum(len(conv['value'].split()) for conv in sample['conversations'])
            cur_len = cur_len if 'image' in sample else -cur_len
            length_list.append(cur_len)
        return length_list

    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
        sources = self.list_data_dict[i]
        if isinstance(i, int):
            sources = [sources]
        assert len(sources) == 1, "Don't know why it is wrapped to a list"  # FIXME
        if 'image' in sources[0]:
            image_file = self.list_data_dict[i]['image']
            image_folder = self.data_args.image_folder
            processor = self.data_args.image_processor
            image_path = os.path.join(image_folder, image_file)
            if not os.path.exists(image_path):
                image_path = image_path[: -3] + "png"
                if not os.path.exists(image_path):
                    image_path = image_path[: -3] + "gif"
            image = Image.open(image_path).convert('RGB')
            if self.data_args.image_aspect_ratio == 'pad':
                def expand2square(pil_img, background_color):
                    width, height = pil_img.size
                    if width == height:
                        return pil_img
                    elif width > height:
                        result = Image.new(pil_img.mode, (width, width), background_color)
                        result.paste(pil_img, (0, (width - height) // 2))
                        return result
                    else:
                        result = Image.new(pil_img.mode, (height, height), background_color)
                        result.paste(pil_img, ((height - width) // 2, 0))
                        return result
                image = expand2square(image, tuple(int(x*255) for x in processor.image_mean))
                image = processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
            else:
                image = processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
            sources = preprocess_multimodal(
                copy.deepcopy([e["conversations"] for e in sources]),
                self.data_args)
        else:
            sources = copy.deepcopy([e["conversations"] for e in sources])
        data_dict = preprocess(
            sources,
            self.tokenizer,
            has_image=('image' in self.list_data_dict[i]), 
            gptneox_tokenizer=self.gptneox_tokenizer,
            coco_caption_sft=self.coco_caption_sft, 
            coco_caption_prompts = self.coco_caption_prompts if self.coco_caption_sft else [])
        if isinstance(i, int):
            data_dict = dict(input_ids=data_dict["input_ids"][0],
                             labels=data_dict["labels"][0])

        # image exist in the data
        if 'image' in self.list_data_dict[i]:
            data_dict['image'] = image
        elif self.data_args.is_multimodal:
            # image does not exist in the data, but the model is multimodal
            crop_size = self.data_args.image_processor.crop_size
            data_dict['image'] = torch.zeros(3, crop_size['height'], crop_size['width'])
        return data_dict


@dataclass
class DataCollatorForSupervisedDataset(object):
    """Collate examples for supervised fine-tuning."""

    tokenizer: transformers.PreTrainedTokenizer

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        input_ids, labels = tuple([instance[key] for instance in instances]
                                  for key in ("input_ids", "labels"))
        input_ids = torch.nn.utils.rnn.pad_sequence(
            input_ids,
            batch_first=True,
            padding_value=self.tokenizer.pad_token_id)
        labels = torch.nn.utils.rnn.pad_sequence(labels,
                                                 batch_first=True,
                                                 padding_value=IGNORE_INDEX)
        input_ids = input_ids[:, :self.tokenizer.model_max_length]
        labels = labels[:, :self.tokenizer.model_max_length]
        batch = dict(
            input_ids=input_ids,
            labels=labels,
            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
        )

        if 'image' in instances[0]:
            images = [instance['image'] for instance in instances]
            if all(x is not None and x.shape == images[0].shape for x in images):
                batch['images'] = torch.stack(images)
            else:
                batch['images'] = images

        return batch


def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer,
                                data_args, coco_caption_sft=False) -> Dict:
    """Make dataset and collator for supervised fine-tuning."""
    train_dataset = LazySupervisedDataset(tokenizer=tokenizer,
                                data_path=data_args.data_path,
                                data_args=data_args, 
                                coco_caption_sft=coco_caption_sft)
    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
    return dict(train_dataset=train_dataset,
                eval_dataset=None,
                data_collator=data_collator)


def train(attn_implementation=None):
    global local_rank

    parser = transformers.HfArgumentParser(
        (ModelArguments, DataArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    # My additions
    if model_args.pretrain_mm_mlp_adapter is None or model_args.pretrain_mm_mlp_adapter.lower() == "none":
        model_args.pretrain_mm_mlp_adapter = None
    local_rank = training_args.local_rank
    # My additions
    # for wandb logging convinence
    # set the `output directory` of wandb 3rd library
    llava_wandb_dir = pathlib.Path(os.getenv("HOME", "/mnt/lustre/lizongshu")) / "outputs" / "LLaVA"
    llava_wandb_dir.mkdir(parents=True, exist_ok=True)
    os.environ["WANDB_DIR"] = str(llava_wandb_dir)

    # set the `project name` of wandb 3rd library
    os.environ["WANDB_PROJECT"] = "improved Mamba-LLaVA experiments"

    # set the `run_name` of wandb 3rd library
    training_args.run_name = pathlib.Path(training_args.output_dir).name
    if training_args.report_to == "wandb":
        rank0_print(f"param `run_name` of wandb lib: {training_args.run_name}")

    compute_dtype = (torch.float16 if training_args.fp16 else (torch.bfloat16 if training_args.bf16 else torch.float32))

    bnb_model_from_pretrained_args = {}
    if training_args.bits in [4, 8]:
        from transformers import BitsAndBytesConfig
        bnb_model_from_pretrained_args.update(dict(
            device_map={"": training_args.device},
            load_in_4bit=training_args.bits == 4,
            load_in_8bit=training_args.bits == 8,
            quantization_config=BitsAndBytesConfig(
                load_in_4bit=training_args.bits == 4,
                load_in_8bit=training_args.bits == 8,
                llm_int8_skip_modules=["mm_projector"],
                llm_int8_threshold=6.0,
                llm_int8_has_fp16_weight=False,
                bnb_4bit_compute_dtype=compute_dtype,
                bnb_4bit_use_double_quant=training_args.double_quant,
                bnb_4bit_quant_type=training_args.quant_type # {'fp4', 'nf4'}
            )
        ))

    # My additions: Noised embeddings finetuning for MLLMs
    if training_args.neftune_alpha is None:
        training_args.neftune_alpha = 0.0

    if model_args.vision_tower is not None:
        if 'mpt' in model_args.model_name_or_path:
            config = transformers.AutoConfig.from_pretrained(model_args.model_name_or_path, trust_remote_code=True)
            config.attn_config['attn_impl'] = training_args.mpt_attn_impl
            model = LlavaMptForCausalLM.from_pretrained(
                model_args.model_name_or_path,
                config=config,
                cache_dir=training_args.cache_dir,
                **bnb_model_from_pretrained_args
            )
        elif 'mamba' in model_args.model_name_or_path:
            config = transformers.AutoConfig.from_pretrained(model_args.model_name_or_path, trust_remote_code=True)
            # My additions
            config.mm_vision_tower = model_args.vision_tower
            config.mm_projector_type = model_args.mm_projector_type
            config.tune_mm_mlp_adapter = model_args.tune_mm_mlp_adapter
            config.mm_vision_select_layer = model_args.mm_vision_select_layer
            config.mm_use_im_start_end = model_args.mm_use_im_start_end
            config.mm_use_im_patch_token = model_args.mm_use_im_patch_token
            config.tokenizer_model_max_length = training_args.model_max_length
            config.mm_patch_merge_type = model_args.mm_patch_merge_type
            config.freeze_mm_mlp_adapter = training_args.freeze_mm_mlp_adapter

            # My additions: Noised embeddings finetuning for MM-Chinese-Mamba
            config.neftune_alpha = model_args.neftune_alpha = training_args.neftune_alpha

            # the `mm_projector_lr` variable is modified, 
            # maybe this will cause the instability in training phase
            config.mm_projector_lr = training_args.mm_projector_lr
            model = LlavaMambaForCausalLM.from_pretrained(
                model_args.model_name_or_path, 
                config=config, 
                cache_dir=training_args.cache_dir, 
                # torch_dtype=(torch.bfloat16 if training_args.bf16 else None), 
                torch_dtype=torch.float,
                ignore_mismatched_sizes=True if model_args.tune_vision_tower else False,
                **bnb_model_from_pretrained_args
            )
            # My additions, for resolution 336 -> 448 reloading pretrained ViT
            if model_args.tune_vision_tower:
                model.get_vision_tower().is_loaded = False
                model.get_vision_tower().load_model()
        else:
            model = LlavaLlamaForCausalLM.from_pretrained(
                model_args.model_name_or_path,
                cache_dir=training_args.cache_dir,
                attn_implementation=attn_implementation,
                torch_dtype=(torch.bfloat16 if training_args.bf16 else None),
                **bnb_model_from_pretrained_args
            )
    else:
        model = transformers.LlamaForCausalLM.from_pretrained(
            model_args.model_name_or_path,
            cache_dir=training_args.cache_dir,
            attn_implementation=attn_implementation,
            torch_dtype=(torch.bfloat16 if training_args.bf16 else None),
            **bnb_model_from_pretrained_args
        )
    model.config.use_cache = False

    if model_args.freeze_backbone:
        if "mamba" in model_args.model_name_or_path:
            model.backbone.requires_grad_(False)
        elif "mpt" in model_args.model_name_or_path:
            model.transformer.requires_grad_(False)
        else:  # LLaMA/Mistral
            model.model.requires_grad_(False)

    if training_args.bits in [4, 8]:
        from peft import prepare_model_for_kbit_training
        model.config.torch_dtype=(torch.float32 if training_args.fp16 else (torch.bfloat16 if training_args.bf16 else torch.float32))
        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=training_args.gradient_checkpointing)

    if training_args.gradient_checkpointing:
        if hasattr(model, "enable_input_require_grads"):
            model.enable_input_require_grads()
        else:
            def make_inputs_require_grad(module, input, output):
                output.requires_grad_(True)
            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)

    if training_args.lora_enable:
        from peft import LoraConfig, get_peft_model
        mamba_target_modules = ["x_proj", "embeddings", "in_proj", "out_proj"]
        lora_config = LoraConfig(
            r=training_args.lora_r,
            lora_alpha=training_args.lora_alpha,
            target_modules=find_all_linear_names(model) if "mamba" not in model_args.model_name_or_path.lower() else mamba_target_modules,
            lora_dropout=training_args.lora_dropout,
            bias=training_args.lora_bias,
            task_type="CAUSAL_LM",
        )
        if training_args.bits == 16:
            # My additions
            if (training_args.bf16) and ('mamba' not in model_args.model_name_or_path):
                model.to(torch.bfloat16)
            if (training_args.fp16) and ('mamba' not in model_args.model_name_or_path):
                model.to(torch.float16)
        rank0_print("Adding LoRA adapters...")
        model = get_peft_model(model, lora_config)

    if ('mpt' in model_args.model_name_or_path.lower()) or ('mamba' in model_args.model_name_or_path.lower()):
        tokenizer = transformers.AutoTokenizer.from_pretrained(
            model_args.model_name_or_path,
            cache_dir=training_args.cache_dir,
            model_max_length=training_args.model_max_length,
            padding_side="right"
        )
    else:
        tokenizer = transformers.AutoTokenizer.from_pretrained(
            model_args.model_name_or_path,
            cache_dir=training_args.cache_dir,
            model_max_length=training_args.model_max_length,
            padding_side="right",
            use_fast=False,
        )

    if model_args.version == "v0":
        if tokenizer.pad_token is None:
            smart_tokenizer_and_embedding_resize(
                special_tokens_dict=dict(pad_token="[PAD]"),
                tokenizer=tokenizer,
                model=model,
            )
    elif model_args.version == "v0.5":
        tokenizer.pad_token = tokenizer.unk_token
    else:
        if 'mamba' not in model_args.model_name_or_path.lower():
            tokenizer.pad_token = tokenizer.unk_token
        else:
            assert model_args.version in ('mpt', 'plain'), "Mamba has the same tokenizer as MPT, i.e. GTPNeoXTokenizer"
        if (model_args.version in ('mpt', 'plain')) and ('mamba' in model_args.model_name_or_path.lower()):
            additional_special_tokens = ["<|im_start|>", "<|im_end|>", "<im_start>", "<im_end>", "<image>", 
                                         "<im_patch>", "<image-placeholder>"]
            smart_tokenizer_and_embedding_resize(special_tokens_dict={"additional_special_tokens": additional_special_tokens}, 
                                                 tokenizer=tokenizer, model=model)
        if model_args.version in conversation_lib.conv_templates:
            conversation_lib.default_conversation = conversation_lib.conv_templates[model_args.version]
            # My additions: change the special end token for coco karpathy training & evaluation
            if (model_args.version == "plain") and ("mamba" in model_args.model_name_or_path.lower()):
                conversation_lib.default_conversation.sep = tokenizer.pad_token
        else:
            conversation_lib.default_conversation = conversation_lib.conv_templates["vicuna_v1"]

    if model_args.vision_tower is not None:
        model.get_model().initialize_vision_modules(
            model_args=model_args,
            fsdp=training_args.fsdp
        )
        
        vision_tower = model.get_vision_tower()
        # My additions
        # When using Mamba as language model, use `torch.amp.autocast` to training by default,
        # so the vision tower in LLaVA (if having a vision tower) should be torch.float
        if 'mamba' not in model_args.model_name_or_path:
            vision_tower.to(dtype=torch.bfloat16 if training_args.bf16 else torch.float16, device=training_args.device)
        else:
            if vision_tower.dtype != torch.float32:
                vision_tower.to(dtype=torch.float32, device=training_args.device)
            else:
                vision_tower.to(device=training_args.device)

        data_args.image_processor = vision_tower.image_processor
        data_args.is_multimodal = True

        model.config.image_aspect_ratio = data_args.image_aspect_ratio
        model.config.tokenizer_padding_side = tokenizer.padding_side
        model.config.tokenizer_model_max_length = tokenizer.model_max_length

        model.config.tune_mm_mlp_adapter = training_args.tune_mm_mlp_adapter = model_args.tune_mm_mlp_adapter
        if model_args.tune_mm_mlp_adapter:
            model.requires_grad_(False)
            for p in model.get_model().mm_projector.parameters():
                p.requires_grad = True

        # My additions, LLDR (layer-wise learning rate decay for CLIP-ViT vision tower)
        training_args.tune_vision_tower = model_args.tune_vision_tower
        training_args.model_name_or_path = model_args.model_name_or_path
        training_args.vision_tower_name_or_path = model_args.vision_tower
        if model_args.tune_vision_tower:
            if training_args.vision_tower_lr is not None:
                assert training_args.vision_tower_lr != 0, "Please provide the non-zero lr for CLIP-ViT when finetuning vision tower"
            else:
                raise ValueError("Please provide a learning rate for CLIP-ViT when finetuning vision tower")
            if training_args.vision_tower_lldr is not None:
                assert training_args.vision_tower_lldr > 0.0 and training_args.vision_tower_lldr <= 1.0, \
                    "the LLDR of vision tower must be greater than 0, smaller than or equal to 1"
            else:
                raise ValueError("Please provide a LLDR for CLIP-ViT when finetuning vision tower")
        setattr(model.config, "tune_vision_tower", model_args.tune_vision_tower)
        setattr(model.config, "vision_tower_lr", training_args.vision_tower_lr)
        setattr(model.config, "vision_tower_lldr", training_args.vision_tower_lldr)
        if model_args.tune_vision_tower:
            model.get_vision_tower().requires_grad_(True)
            for p in model.get_vision_tower().parameters():
                p.requires_grad = True
            # My additions, When training vision encoder, some layers at the last are unused parameters
            # i.e. their requires_grad is true, but they don't participate the loss caculation
            training_args.ddp_find_unused_parameters = True
            training_args.gradient_checkpointing_kwargs = {"use_reentrant": False}

        model.config.freeze_mm_mlp_adapter = training_args.freeze_mm_mlp_adapter
        if training_args.freeze_mm_mlp_adapter:
            for p in model.get_model().mm_projector.parameters():
                p.requires_grad = False

        if training_args.bits in [4, 8]:
            model.get_model().mm_projector.to(dtype=compute_dtype, device=training_args.device)

        model.config.mm_use_im_start_end = data_args.mm_use_im_start_end = model_args.mm_use_im_start_end
        model.config.mm_projector_lr = training_args.mm_projector_lr
        training_args.use_im_start_end = model_args.mm_use_im_start_end
        model.config.mm_use_im_patch_token = model_args.mm_use_im_patch_token
        model.initialize_vision_tokenizer(model_args, tokenizer=tokenizer)
        # My additions
        native_pinpoints = model.config.image_grid_pinpoints
        new_resolution = int(model_args.vision_tower.split("/")[-1].strip().strip("/").strip().split("-")[-1])
        if native_pinpoints[0][0] != new_resolution:
            mult = new_resolution / native_pinpoints[0][0]
            tmp_pinpoints = [pinpoint[:] for pinpoint in native_pinpoints]
            new_pinpoints = [[int(point * mult) for point in pinpoint] for pinpoint in tmp_pinpoints]
            model.config.image_grid_pinpoints = new_pinpoints

    # My additions
    model.config.data_path = training_args.data_path = data_args.data_path
    model.config.image_folder = training_args.image_folder = data_args.image_folder

    if training_args.bits in [4, 8]:
        from peft.tuners.lora import LoraLayer
        for name, module in model.named_modules():
            if isinstance(module, LoraLayer):
                if training_args.bf16:
                    module = module.to(torch.bfloat16)
            if 'norm' in name:
                module = module.to(torch.float32)
            if 'lm_head' in name or 'embed_tokens' in name or 'embeddings' in name:
                if hasattr(module, 'weight'):
                    if training_args.bf16 and module.weight.dtype == torch.float32:
                        module = module.to(torch.bfloat16)

    data_module = make_supervised_data_module(tokenizer=tokenizer,
                                              data_args=data_args, 
                                              coco_caption_sft=training_args.coco_caption_sft)
    trainer = LLaVATrainer(model=model,
                    tokenizer=tokenizer,
                    args=training_args,
                    **data_module)

    if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):
        trainer.train(resume_from_checkpoint=True)
    else:
        trainer.train()
    trainer.save_state()

    model.config.use_cache = True

    if training_args.lora_enable:
        state_dict = get_peft_state_maybe_zero_3(
            model.named_parameters(), training_args.lora_bias
        )
        non_lora_state_dict = get_peft_state_non_lora_maybe_zero_3(
            model.named_parameters()
        )
        if training_args.local_rank == 0 or training_args.local_rank == -1:
            model.config.save_pretrained(training_args.output_dir)
            model.save_pretrained(training_args.output_dir, state_dict=state_dict)
            torch.save(non_lora_state_dict, os.path.join(training_args.output_dir, 'non_lora_trainables.bin'))
    else:
        safe_save_model_for_hf_trainer(trainer=trainer,
                                       output_dir=training_args.output_dir)


if __name__ == "__main__":
    train()
